#include <complex.h>
#include <stdio.h>
#include <math.h>
#include <jack/jack.h>
#include <fftw3.h>
#include <errno.h>

#include "main.h"

#define N 512

// gain and phase correction variables
double gain[MAX_CHANNELS][N], phase[MAX_CHANNELS][N];

inline void GoFrequencyDomain(jack_default_audio_sample_t *x, jack_default_audio_sample_t *y);

int fit_curve(int n, double yi[]);

extern fftw_complex *fft_in, *fft_out, *fft_in_inv, *fft_out_inv;

int ComputeIQBalance(int channel, jack_default_audio_sample_t *x, jack_default_audio_sample_t *y) {
    int bin;

    // move to the frequency domain if not there already
    GoFrequencyDomain(x, y);

    double average = 0;
    // scan for bin with magnitude 30 db above average - scan once to get average
    for (bin = 0; bin < N; bin++) {
        average += cabs(fft_out[bin]);
    }
    average /= N;

    // scan for bin with magnitude 20 db above average
    for (bin = 0; bin < N; bin++) {
        if (cabs(fft_out[bin]) < average * 10.0) {
            continue;
        }

        // use Alex, VE3NEA's algorithm to compute the correction factors for this signal and
        // image pair
        complex signal = fft_out[bin];
        complex image = fft_out[N - bin];
        complex Z = signal * image;
        double power = creal(signal) * creal(signal) + cimag(signal) * cimag(signal) +
            creal(image) * creal(image) + cimag(image) * cimag(image);
        complex Zprime = creal(Z) / power + I * cimag(Z) / power;
        double Temp = sqrt(1.0 - 4.0 * creal(Zprime) * creal(Zprime));

        double Phase = asin(2.0 * cimag(Zprime) / Temp);
        double Gain = Temp / (1.0 - 2.0 * creal(Zprime));

        gain[channel][bin] = Gain;
        phase[channel][bin] = Phase;

        printf("%4d ", bin);
    }
    printf("\n");

    // now fit curves to the gain and phase correction data.
    //fit_curve(N, gain[channel]);
    //fit_curve(N, phase[channel]);

    // return remaining in the frequency domain
    return 0;
}

int BalanceIQ(int channel, jack_default_audio_sample_t *x, jack_default_audio_sample_t *y) {
    int bin;

    // move to the frequency domain if not there already
    GoFrequencyDomain(x, y);

    // apply the gain and balance curves to the frequency domain signals
    for (bin = 0; bin < N; bin++) {
        double x_r = creal(fft_out[bin]);
        double x_i = cimag(fft_out[bin]);
        double x_r_sym = creal(fft_out[N - bin]);
        double x_i_sym = cimag(fft_out[N - bin]);
        double E = cos(phase[channel][bin]) / gain[channel][bin] - 1.0;
        double P = -sin(phase[channel][bin]) / gain[channel][bin];
        double complex X_i = 0.5 * (x_r + x_r_sym) + I * 0.5 * (x_i - x_i_sym);
        fft_in_inv[bin] = (x_r + I * x_i) + ((E + I * P) * X_i);
    }

    // return remaining in the frequency domain
    return 0;
}

int DumpIQPoints(FILE *file, int channel) {
    int i, bin, dump_bytes = 0;
    for (bin = -N / 2; bin < N / 2; bin++) {
        i = bin < 0 ? N + bin : bin;
        if (gain[channel][i] == 0)
            continue;
        dump_bytes += fprintf(file, "%d %f %f\n", i, gain[channel][i], phase[channel][i]);
    }
    return dump_bytes;
}

#include <gsl/gsl_multifit.h>
 
int fit_curve(int n, double yi[]) {
    int i;
    double xi, chisq;
    gsl_matrix *X, *cov;
    gsl_vector *y, *c;

    X = gsl_matrix_alloc (n, 3);
    y = gsl_vector_alloc (n);
     
    c = gsl_vector_alloc (3);
    cov = gsl_matrix_alloc (3, 3);
     
    for (i = 0; i < n; i++) {
        if (yi[i] == 0)
            continue;
        xi = i >= N / 2 ? N - i : i;
           
        gsl_matrix_set (X, i, 0, 1.0);
        gsl_matrix_set (X, i, 1, xi);
        gsl_matrix_set (X, i, 2, xi*xi);
        
        gsl_vector_set (y, i, yi[i]);
    }
     
    {
        gsl_multifit_linear_workspace * work = gsl_multifit_linear_alloc (n, 3);
        gsl_multifit_linear (X, y, c, cov, &chisq, work);
        gsl_multifit_linear_free (work);
    }
     
#define C(i) (gsl_vector_get(c,(i)))
#define COV(i,j) (gsl_matrix_get(cov,(i),(j)))
     
#if 0
    {
        printf ("# best fit: Y = %g + %g X + %g X^2\n", C(0), C(1), C(2));
     
        printf ("# covariance matrix:\n");
        printf ("[ %+.5e, %+.5e, %+.5e  \n", COV(0,0), COV(0,1), COV(0,2));
        printf ("  %+.5e, %+.5e, %+.5e  \n", COV(1,0), COV(1,1), COV(1,2));
        printf ("  %+.5e, %+.5e, %+.5e ]\n", COV(2,0), COV(2,1), COV(2,2));
        printf ("# chisq = %g\n", chisq);
    }
#endif

    // fill in the sparse array
    for (i = 0; i < N; i++) {
        xi = i >= N / 2 ? N - i : i;
        yi[i] = C(0) + C(1) * xi + C(2) * xi * xi;
    }
     
    gsl_matrix_free (X);
    gsl_vector_free (y);
    gsl_vector_free (c);
    gsl_matrix_free (cov);

    return 0;
}
